import os
import re
from contextlib import contextmanager

import dotenv
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_deepseek import ChatDeepSeek
from langchain_ollama import ChatOllama
from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker


class ReportGenerator:
    def __init__(self):
        """初始化"""
        dotenv.load_dotenv()

        self.extractor = ChatOllama(
            model="qwen3:32b",
            base_url=os.getenv("OLLAMA_BASE_URL"),
            temperature=0.4,
            top_p=0.3,
            top_k=10,
            seed=0,
            reasoning=False,
            keep_alive='5m',
            format='json'
        )

        self.summary = ChatOllama(
            model="deepseek-r1:70b",
            base_url=os.getenv("OLLAMA_BASE_URL"),
            temperature=0.2,
            top_p=0.1,
            top_k=21,
            keep_alive='5m',
            seed=0
        )

        self.deepseek_coder = ChatDeepSeek(
            model="deepseek-chat",
            temperature=0.4,
            top_p=0.4
        )

        self.desensitizer = ChatOllama(
            model="qwen3:32b",
            base_url=os.getenv("OLLAMA_BASE_URL"),
            temperature=0.4,
            top_p=0.5,
            top_k=20,
            seed=0,
            keep_alive='5m'
        )

        # 数据库连接
        user = os.getenv('POSTGRES_USER')
        password = os.getenv('POSTGRES_PASSWORD')
        host = os.getenv('POSTGRES_HOST')
        port = os.getenv('POSTGRES_PORT')
        database = os.getenv('POSTGRES_DB')

        self.engine = create_engine(
            f"postgresql+psycopg2://{user}:{password}@{host}:{port}/{database}"
        )
        self.session_maker = sessionmaker(bind=self.engine)

    @contextmanager
    def get_db_session(self):
        """数据库会话管理器"""
        session = self.session_maker()
        try:
            yield session
            session.commit()
        except Exception as e:
            session.rollback()
            raise e
        finally:
            session.close()

    @staticmethod
    def load_system_prompt(file_path):
        """加载系统提示词"""
        try:
            with open(file_path, "r", encoding="utf-8") as f:
                return f.read()
        except FileNotFoundError:
            raise FileNotFoundError(f"系统提示文件 {file_path} 未找到")
        except Exception as e:
            raise Exception(f"读取系统提示文件失败: {str(e)}")

    def extract_key_info(self, user_input):
        """提取用户输入中的关键信息"""
        user_prompt = f"""根据用户输入：{user_input}，提取用户语句中的关键信息，例如供电局名称，线路名称，线路电压等信息，提取关键信息。
        """

        messages = [
            SystemMessage(content="不需要进行分点与列表，不需要多余的解释与思考。"),
            HumanMessage(content=user_prompt)
        ]

        return self.remove_think_tag(self.extractor.invoke(messages).text())

    def generate_sql_statements(self, station_info):
        """根据关键信息生成SQL语句"""
        sql_system_prompt = self.load_system_prompt("../system_prompts/system_prompt_sql.txt")

        sql_messages = [
            SystemMessage(content=sql_system_prompt),
            HumanMessage(content="线路信息如下：\n" + station_info),
        ]

        sql_response = self.deepseek_coder.invoke(sql_messages)
        return sql_response.content.split("\n\n")

    def execute_sql_queries(self, sql_statements):
        """执行SQL查询"""
        results = []

        with self.get_db_session() as session:
            for sql_statement in sql_statements:
                if sql_statement.strip():  # 跳过空语句
                    try:
                        result = session.execute(text(sql_statement))
                        rows = [dict(row._mapping) for row in result]
                        results.append({"records": rows})
                    except Exception as e:
                        print(f"执行SQL语句失败: {sql_statement}, 错误: {str(e)}")
                        results.append(None)

        return results


    def generate_report(self, station_info, query_results):
        """生成报告"""
        report_system_prompt = self.load_system_prompt("../system_prompts/system_prompt_report.txt")

        prt = [
            SystemMessage(content=report_system_prompt),
            HumanMessage(content=f'''【参数信息】\n{station_info}\n【实际数据】\n{query_results}''')
        ]
        report_response = self.summary.invoke(prt)

        return self.remove_think_tag(report_response.content)

    def desensitize_report(self, report_text):
        """报告脱敏处理"""
        desensitization_system_prompt = self.load_system_prompt("../system_prompts/system_prompt_desensitize.txt")

        prt = [
            SystemMessage(content=desensitization_system_prompt),
            HumanMessage(content=report_text)
        ]
        final_response = self.desensitizer.invoke(prt)

        return self.remove_think_tag(final_response.text())

    @staticmethod
    def remove_think_tag(text: str) -> str:
        return re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)

    def analyze_power_lines(self, user_input):
        """主分析流程"""
        try:
            # 1. 提取关键信息
            station_info = self.extract_key_info(user_input)

            # 2. 生成SQL语句
            sql_statements = self.generate_sql_statements(user_input)

            # 3. 执行SQL查询
            query_results = self.execute_sql_queries(sql_statements)

            query_results_str = '\n'.join([str(item) for item in query_results])

            # 4. 生成报告
            report_text = self.generate_report(station_info, query_results_str)

            # 5. 报告脱敏
            final_report = self.desensitize_report(report_text)

            return final_report

        except Exception as e:
            print(f"分析过程中发生错误: {str(e)}")
            raise e

# 使用示例
if __name__ == "__main__":
    analyzer = ReportGenerator()
    user_input = "昆明供电局220kV草水双回轮停、220kV水城变220kVⅠ、Ⅱ母轮停、220kV#2主变停电"

    try:
        result = analyzer.analyze_power_lines(user_input)
        print(result)
    except Exception as e:
        print(f"分析失败: {str(e)}")
